# -*- coding: utf-8 -*-
"""
本脚本用于整合同项目 fundus-reasoner-datapreprocess 产出的 Alpaca 数据，
按“阶段”合并多个已生成的数据集为单一 SFT 训练集，并在目标目录写出
dataset_info.json 以登记该阶段数据集的元信息。

用法：
  - 直接编辑下方全局变量 `stage_1_set`，把你希望合并的“数据集名称或文件名”填进去。
    * 名称：应能在源目录 alpaca-data_dir 下的 dataset_info.json 中找到对应条目
    * 文件名：直接指向源目录中的 .json 文件（如 "med_pub_simpleQA_train.json"）
  - 运行本脚本即可在 `tar_dir` 下生成：
      - stage1_sft_med_qa.json          # 合并后的训练数据
      - dataset_info.json               # 注册该阶段数据集的元信息（可多次合并追加）

注意：
  - 仅做简单的去重（基于 (instruction, output, images[0]) 三元组哈希）
  - 不对样本字段做变换，假定与既有 Alpaca 列规范一致：
        instruction / input / output / images
"""

import os
import json
import random
from typing import List, Dict, Any, Tuple

# 可选：用 loguru 打印更友好的日志；没有的话退化为 print
try:
    from loguru import logger
    _HAS_LOGURU = True
except Exception:
    _HAS_LOGURU = False
    class _P:
        def info(self, *a, **k): print("[INFO]", *a)
        def success(self, *a, **k): print("[OK]", *a)
        def warning(self, *a, **k): print("[WARN]", *a)
        def error(self, *a, **k): print("[ERR ]", *a)
    logger = _P()

# ===================== 全局配置 =====================
# 源目录（训练/测试构建脚本已经把各 json 放在这里）
alpaca_data_dir = "/home/zhangpinglu/data0/gy/Dataset/fundusreasoner/Alpaca_data/"
# 目标目录（阶段数据与其 dataset_info.json 会写到这里）
tar_dir = "/home/zhangpinglu/data0/gy/Dataset/fundusreasoner/dataset/"
os.makedirs(tar_dir, exist_ok=True)

# —— 在这里列出要并入 Stage 1 的数据集（名称或文件名均可）——
stage_1_set: List[str] = [
    "med_pub_simpleQA_train",
    "med_pub_coor_qa_train",
    "aier_coor_qa_train"
]

# ===================== 基础IO与工具 =====================

def save_json(data: Any, path: str) -> None:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    logger.success(f"已保存：{path}")

def read_json(path: str) -> Any:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

def merge_dataset_info(alpaca_dir: str, info_dict: Dict[str, Any]) -> None:
    """将 info_dict 合并写入 alpaca_dir/dataset_info.json"""
    info_path = os.path.join(alpaca_dir, "dataset_info.json")
    if os.path.exists(info_path):
        with open(info_path, "r", encoding="utf-8") as f:
            all_info = json.load(f)
    else:
        all_info = {}
    all_info.update(info_dict)
    save_json(all_info, info_path)
    logger.success("dataset_info.json 已更新。")

def _load_source_dataset_info(src_dir: str) -> Dict[str, Any]:
    """读取源目录的 dataset_info.json（若不存在，返回空 dict）"""
    info_path = os.path.join(src_dir, "dataset_info.json")
    if os.path.exists(info_path):
        try:
            return read_json(info_path)
        except Exception as e:
            raise RuntimeError(f"[读取源 dataset_info.json 失败] path={info_path}, error={e}")
    return {}

def _resolve_source_file(name_or_file: str, src_info: Dict[str, Any]) -> Tuple[str, str]:
    """
    将“数据集名称或文件名”解析为源文件绝对路径与展示名。
    返回：(abs_json_path, display_name)
    解析策略：
      1) 若字符串以 .json 结尾，视作文件名 => 直接在 alpaca_data_dir 下寻找
      2) 否则视作 “数据集名称”，在源 dataset_info.json 中查找 file_name
    """
    if name_or_file.lower().endswith(".json"):
        abs_path = os.path.join(alpaca_data_dir, name_or_file)
        if not os.path.exists(abs_path):
            raise FileNotFoundError(f"[未找到源文件] {abs_path}")
        return abs_path, os.path.splitext(os.path.basename(name_or_file))[0]

    if name_or_file not in src_info:
        raise KeyError(f"[源 dataset_info 缺少条目] name='{name_or_file}'")
    file_name = src_info[name_or_file].get("file_name")
    if not file_name:
        raise KeyError(f"[源 dataset_info 条目缺少 file_name] name='{name_or_file}' -> {src_info[name_or_file]}")
    abs_path = os.path.join(alpaca_data_dir, file_name)
    if not os.path.exists(abs_path):
        raise FileNotFoundError(f"[未找到源文件] name='{name_or_file}', file='{abs_path}'")
    return abs_path, name_or_file

def _load_alpaca_list(path: str) -> List[Dict[str, Any]]:
    """读取 Alpaca 列表并做最基本检查"""
    data = read_json(path)
    if not isinstance(data, list):
        raise ValueError(f"[Alpaca 文件格式错误] 期望 list，实际 {type(data).__name__} at {path}")
    return data

def _uniq_by_triplet(items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    基于 (instruction, output, images[0]) 做简单去重。
    若 images 为空或不是 list，则退化仅用 instruction/output。
    """
    seen = set()
    out = []
    for obj in items:
        instr = obj.get("instruction", "")
        outp = obj.get("output", "")
        imgs = obj.get("images", [])
        key = (instr, outp, imgs[0] if isinstance(imgs, list) and len(imgs) > 0 else "")
        if key in seen:
            print(f"去重样本：{key}") # 一般不会触发，这里防一手clear
            
            continue
        seen.add(key)
        out.append(obj)
    return out

# ===================== 阶段合并实现 =====================

def stage1_sft_med_qa(
    out_name: str = "stage1_sft_med_qa.json",
    shuffle_seed: int = 42
) -> None:
    """
    把 stage_1_set 列出的数据集合并成一个训练集，写入 tar_dir/out_name，
    并在 tar_dir/dataset_info.json 中登记该集合（仅包含 file_name + columns 字段）。
    """
    if not stage_1_set:
        raise ValueError("stage_1_set 为空，请在脚本顶部填入要合并的数据集/文件名。")

    # 读取源目录 dataset_info（帮助把“数据集名称”解析为文件名）
    src_info = _load_source_dataset_info(alpaca_data_dir)

    # 解析所有源文件路径
    src_files: List[Tuple[str, str]] = []  # (abs_path, display_name)
    for name in stage_1_set:
        abs_path, display = _resolve_source_file(name, src_info)
        src_files.append((abs_path, display))
        logger.info(f"纳入合并：{display} -> {abs_path}")

    # 读取并拼接
    merged: List[Dict[str, Any]] = []
    per_file_counts = {}
    for abs_path, display in src_files:
        arr = _load_alpaca_list(abs_path)
        merged.extend(arr)
        per_file_counts[display] = len(arr)

    # 去重与打乱
    before_n = len(merged)
    merged = _uniq_by_triplet(merged)
    after_n = len(merged)
    random.Random(shuffle_seed).shuffle(merged)

    logger.info(f"合并完成：源总数={before_n}，去重后={after_n}。各源计数：{per_file_counts}")

    # 写主数据
    out_path = os.path.join(tar_dir, out_name)
    save_json(merged, out_path)

    # === 完成 TODO：将 info 条目改成精简版形态（仿照示例） ===
    info_key = os.path.splitext(out_name)[0]  # 例如 "stage1_sft_med_qa"
    info_entry = {
        info_key: {
            "file_name": out_name,
            "columns": {
                "prompt": "instruction",
                "query": "input",
                "response": "output",
                "images": "images"
            }
        }
    }
    merge_dataset_info(tar_dir, info_entry)
    logger.success(f"✅ Stage 1 数据集完成：{out_path}（已以精简信息写入 dataset_info.json）")
# ===================== Mini 训练集采样 =====================

def make_mini_subset(
    source_file: str = "stage1_sft_med_qa.json",
    ratio: float = 0.01,
    seed: int = 42
) -> None:
    """
    从 tar_dir 下已存在的阶段数据中随机采样 ratio 比例样本，
    生成一个 mini 训练集，用于快速调试。
    """
    src_path = os.path.join(tar_dir, source_file)
    if not os.path.exists(src_path):
        raise FileNotFoundError(f"[mini_train] 源文件不存在：{src_path}")

    # 读取数据
    data = read_json(src_path)
    if not isinstance(data, list):
        raise ValueError(f"[mini_train] 文件格式错误，应为 list，实际是 {type(data).__name__}")

    # 采样
    random.seed(seed)
    n_total = len(data)
    n_sample = max(1, int(n_total * ratio))
    sampled = random.sample(data, n_sample)

    # 写出 mini 文件
    mini_name = os.path.splitext(source_file)[0] + "_mini_train.json"
    mini_path = os.path.join(tar_dir, mini_name)
    save_json(sampled, mini_path)
    logger.success(f"✅ mini 训练集生成完成：{mini_path}（采样 {n_sample}/{n_total} 条）")

    # 写入 dataset_info.json
    info_key = os.path.splitext(mini_name)[0]
    info_entry = {
        info_key: {
            "file_name": mini_name,
            "columns": {
                "prompt": "instruction",
                "query": "input",
                "response": "output",
                "images": "images"
            }
        }
    }
    merge_dataset_info(tar_dir, info_entry)
    logger.success("✅ dataset_info.json 已登记 mini 训练集。")

# ===================== 主入口 =====================

def main():
    # stage1_sft_med_qa()
    make_mini_subset("stage1_sft_med_qa.json")
if __name__=='__main__':
    main()
